#!/usr/bin/env python3
"""
单个 Chirp 快时间 FFT（精简版）

仅保留：
- 对指定 chirp 的快时间 FFT
- 仅加窗（window）与 NFFT 选择
- 绘制并返回双边谱（±频率）与对应的正负距离轴（若提供扫频率率）

已删除：相干积累、去直流、去趋势、相位对齐、慢时间处理等所有其它计算流程。

用法示例
# 单个 chirp 分析（双边谱 + 正负距离）
python -m src.single_chirp_fft --matrix data/your_chirp_matrix.npy --fs 1e6 --window hann --nfft 8192 --chirp-index 0 --save --show
"""

from __future__ import annotations

import argparse
import glob
import os
from pathlib import Path
from typing import Dict, Optional

import numpy as np
import matplotlib.pyplot as plt
try:
    # 尝试使用项目内配置管理器
    from src.config_manager import ConfigManager  # type: ignore
except Exception:
    ConfigManager = None  # 允许无依赖运行

# 设置字体支持
def setup_font():
    """设置支持中文的字体"""
    import matplotlib.font_manager as fm
    
    # 首先尝试常见的安全字体
    safe_fonts = ['Liberation Sans', 'Arial', 'DejaVu Sans', 'Helvetica', 'sans-serif']
    available_fonts = [f.name for f in fm.fontManager.ttflist]
    
    # 找到第一个可用的安全字体
    selected_font = None
    for font in safe_fonts:
        if font in available_fonts:
            selected_font = font
            break
    
    if selected_font:
        plt.rcParams['font.sans-serif'] = [selected_font]
        plt.rcParams['axes.unicode_minus'] = False
    
    # 直接返回 False，强制使用英文标签
    return False

# 初始化字体设置
has_chinese_font = setup_font()

class _Logger:
    def info(self, m: str):
        print(m)
    def warning(self, m: str):
        print(f"[WARN] {m}")
    def error(self, m: str):
        print(f"[ERR] {m}")


def _make_window(name: str, n: int) -> np.ndarray:
    """生成窗函数"""
    name = (name or "").lower()
    if name in ("hann", "hanning"):
        return np.hanning(n)
    if name == "hamming":
        return np.hamming(n)
    if name == "blackman":
        return np.blackman(n)
    if name in ("rect", "rectangular", "boxcar", "none", ""):
        return np.ones(n)
    raise ValueError(f"未知窗函数: {name}")


def _next_pow2(n: int) -> int:
    """计算下一个2的幂"""
    return 1 if n <= 1 else 1 << (int(n - 1).bit_length())


def _find_latest_chirp_matrix(data_dir: str = "./data") -> Path:
    """
    自动扫描 data 目录，找到最新的 chirp_matrix_*.npy 文件（排除 axes 文件）
    
    Args:
        data_dir: 数据目录路径
        
    Returns:
        最新文件的 Path 对象
        
    Raises:
        FileNotFoundError: 如果未找到任何 chirp_matrix 文件
    """
    pattern = os.path.join(data_dir, "**", "chirp_matrix_*.npy")
    all_matches = glob.glob(pattern, recursive=True)
    
    # 过滤掉 axes 文件
    matches = [f for f in all_matches if "axes" not in os.path.basename(f)]
    
    if not matches:
        raise FileNotFoundError(f"在 {data_dir} 中未找到任何有效的 chirp_matrix_*.npy 文件（已排除 axes 文件）")
    
    # 按修改时间排序，选择最新的
    latest_file = max(matches, key=os.path.getmtime)
    return Path(latest_file)


def range_axis_from_freq(freq_hz: np.ndarray, sweep_rate_hz_per_s: Optional[float]) -> Optional[np.ndarray]:
    """
    给定 beat 频率与扫频率率 k（Hz/s），计算距离轴: R = c * f_b / (2k)
    支持负频率，可以显示负距离。
    未提供 k 时返回 None。
    """
    if sweep_rate_hz_per_s is None or sweep_rate_hz_per_s <= 0:
        return None
    c = 3.0e8
    return c * freq_hz / (2.0 * float(sweep_rate_hz_per_s))


# 相干积累等功能已移除


def analyze_single_chirp(
    chirp_matrix: np.ndarray,
    chirp_index: int,
    fs: float,
    sweep_rate_hz_per_s: Optional[float] = None,
    window: str = "hann",
    nfft: Optional[int] = None,
    save_plot: bool = False,
    show_plot: bool = False,
) -> Dict[str, np.ndarray]:
    """
    对快慢时间矩阵中选定 chirp 做 FFT，并可选保存图像。

    Returns: 包含 'freq_hz'、'power_db'、可选 'range_m' 与原始时间波形。
    """
    lg = _Logger()

    if chirp_matrix.ndim != 2:
        raise ValueError("chirp_matrix 应为二维 [慢时间 × 快时间]")
    n_chirps, n_fast = chirp_matrix.shape
    if not (-n_chirps <= chirp_index < n_chirps):
        raise IndexError(f"chirp_index 越界: {chirp_index}, 可选范围 [0, {n_chirps-1}]")

    chirp = chirp_matrix[chirp_index]
    lg.info(f"选取 Chirp #{chirp_index}: 快时间点数={n_fast}")

    # FFT 分析
    if chirp.ndim != 1:
        raise ValueError("chirp 必须为一维序列")
    x = chirp.astype(np.complex128, copy=False)
    
    n = len(x)
    nfft = int(nfft) if nfft and nfft > 0 else _next_pow2(n)

    w = _make_window(window, n)
    xw = x * w

    if nfft > n:
        buf = np.zeros(nfft, dtype=np.complex128)
        buf[:n] = xw
        xw = buf
    else:
        xw = xw[:nfft]

    X = np.fft.fft(xw, n=nfft)
    freq = np.fft.fftfreq(nfft, d=1.0 / fs)

    # 使用 fftshift 让频率单调升序（-fs/2..+fs/2），避免“封闭曲线”
    X_plot = np.fft.fftshift(X)
    freq_plot = np.fft.fftshift(freq)

    # dBFS 标度（幅度归一到 NFFT）
    power_full = (np.abs(X_plot) ** 2) / (nfft ** 2)
    power_db_full = 10.0 * np.log10(power_full + 1e-20)

    # 计算完整频率轴对应的距离轴（支持负距离）
    range_m = range_axis_from_freq(freq_plot, sweep_rate_hz_per_s)

    # 绘图
    if save_plot or show_plot:
        fig = plt.figure(figsize=(10, 8))
        
    # 频谱图（上）- 双边谱（已排序）
        ax1 = fig.add_subplot(2, 1, 1)
        ax1.plot(freq_plot, power_db_full, lw=1.2)
        ax1.set_xlabel("Frequency (Hz)")
        ax1.set_title(f"Fast-time FFT Spectrum (Chirp #{chirp_index}) - Full Spectrum")
        ax1.set_ylabel("Power (dBFS)")
        ax1.grid(True, alpha=0.3)
        
        # 距离谱图（下）- 正负距离（已排序）
        ax2 = fig.add_subplot(2, 1, 2)
        if range_m is not None:
            ax2.plot(range_m, power_db_full, lw=1.2)
            ax2.set_xlabel("Range (m)")
            ax2.set_title(f"Range Spectrum (Chirp #{chirp_index}) - Full Range")
            ax2.set_ylabel("Power (dBFS)")
        else:
            ax2.plot(freq_plot, power_db_full, lw=1.2)
            ax2.set_xlabel("Frequency (Hz)")
            ax2.set_title(f"Fast-time FFT Spectrum (Chirp #{chirp_index}) - Full Spectrum")
            ax2.set_ylabel("Power (dBFS)")
        ax2.grid(True, alpha=0.3)
        
        fig.tight_layout()

        if save_plot:
            outdir = Path("./data")
            outdir.mkdir(parents=True, exist_ok=True)
            out = outdir / f"single_chirp_{chirp_index}.png"
            fig.savefig(out, dpi=300, bbox_inches="tight")
        if show_plot:
            plt.show()
        plt.close(fig)

    return {
        "freq_hz": freq_plot,  # 已排序（-fs/2..+fs/2）或正频率
        "power_db": power_db_full,
        "range_m": range_m if range_m is not None else np.array([]),
        "time_waveform": chirp,
    }


def main():
    ap = argparse.ArgumentParser(description="对快慢时间矩阵的指定 chirp 做 FFT（双边谱/正负距离）")
    ap.add_argument("--matrix", type=str, help="输入 chirp 矩阵 .npy，或 glob 模式；未指定时自动选择最新文件")
    ap.add_argument("--chirp-index", type=int, default=0, help="选择的 chirp 行索引")
    ap.add_argument("--fs", type=float, required=False, help="采样率 Hz")
    ap.add_argument("--slope", type=float, help="FMCW 扫频率率 k=Hz/s，用于距离轴换算")
    ap.add_argument("--nfft", type=int, help="FFT 点数；默认使用下一次幂")
    ap.add_argument("--window", type=str, default="blackman", help="窗函数: hann/hamming/blackman/rect")
    ap.add_argument("--save", action="store_true", help="保存图像到 data 目录或会话目录")
    ap.add_argument("--show", action="store_true", help="弹出图窗显示频谱图")
    args = ap.parse_args()
    lg = _Logger()
    fs_value: Optional[float] = float(args.fs) if args.fs is not None else None

    # 自动加载参数（来自配置文件）
    data_dir = "./data"
    if ConfigManager is not None:
        try:
            cm = ConfigManager("config/config.json")
            cm.load_config()
            # 自动加载采样率与扫频率率
            if fs_value is None:
                fs_cfg = cm.get("instruments.fsw.sample_rate", None)
                if fs_cfg is not None:
                    fs_value = float(fs_cfg)
                    lg.info(f"从配置加载采样率 fs={fs_value}")
            if args.slope is None:
                slope_cfg = cm.get("fmcw_signal.sweep_rate", None)
                if slope_cfg is not None:
                    args.slope = float(slope_cfg)
                    lg.info(f"从配置加载扫频率率 k={args.slope}")
            # 自动加载 FFT 参数
            if args.nfft is None:
                nfft_cfg = cm.get("fft_analysis.fft_points", None)
                if nfft_cfg is not None:
                    args.nfft = int(nfft_cfg)
                    lg.info(f"从配置加载 NFFT={args.nfft}")
            if args.window == "hann":
                win_cfg = cm.get("fft_analysis.window_function", None)
                if isinstance(win_cfg, str) and win_cfg:
                    args.window = win_cfg
                    lg.info(f"从配置加载窗函数: {args.window}")
            # 数据目录（用于自动选择最新文件）
            data_dir = cm.get("output.data_directory", data_dir) or data_dir
        except Exception as e:
            lg.warning(f"无法自动加载配置参数: {e}")

    # 解析矩阵文件路径
    if args.matrix is None:
        # 自动选择最新文件
        try:
            mpath = _find_latest_chirp_matrix(data_dir=data_dir)
            lg.info(f"自动选择最新文件: {mpath}")
        except FileNotFoundError as e:
            lg.error(str(e))
            return
    else:
        mpath = Path(args.matrix)
        if not mpath.exists():
            raise FileNotFoundError(f"未找到矩阵文件: {args.matrix}")
        lg.info(f"使用指定文件: {mpath}")
    
    lg.info(f"加载矩阵: {mpath}")
    mat = np.load(mpath)
    assert mat.ndim == 2, "矩阵应为二维 [慢时间 × 快时间]"
    
    if fs_value is None:
        raise ValueError("未提供采样率 --fs。为简化脚本，已移除配置读取，请显式指定 --fs。")

    res = analyze_single_chirp(
        chirp_matrix=mat,
        chirp_index=int(args.chirp_index),
        fs=float(fs_value),
        sweep_rate_hz_per_s=(float(args.slope) if args.slope else None),
        window=args.window,
        nfft=(int(args.nfft) if args.nfft else None),
        save_plot=bool(args.save),
        show_plot=bool(args.show),
    )

    # 基本结果打印
    freq = res["freq_hz"]
    pdb = res["power_db"]
    peak_idx = int(np.argmax(pdb))
    lg.info(
        f"Nfreq={len(freq)}, 峰值频率={freq[peak_idx]:.3f} Hz, 峰值功率={pdb[peak_idx]:.2f} dBFS"
    )
    if res.get("range_m", np.array([])).size > 0:
        rng = res["range_m"][peak_idx]
        lg.info(f"对应距离峰值 ≈ {rng:.3f} m")


if __name__ == "__main__":
    main()
